import sqlite3
import json
from pyecharts import WordCloud
from pyecharts import Sankey

if __name__ == '__main__':

    f_data = open('data.txt','w')
    f_links = open('links.txt','w')
    conn = sqlite3.connect('test.db')
    cursor = conn.cursor()

    # cursor.execute('create table bankuai (stock_id varchar(20), stock_name varchar(20),yyb_id varchar(20),yyb_name varchar(32),bk_id varchar(20),bk_name varchar(20))')

    # 取营业部id
    cursor.execute('select yyb_id,yyb_name from bankuai  group by yyb_id')

    banks = cursor.fetchall()

    print(banks,'----------------------\r\n')

    name = []
    re = []
    for bk in banks:
        # bk
        if(bk[0]):

            print(bk[0])


            # sql查询
            sql = "select bk_name as tartget ,count(bk_id) as value  from bankuai where yyb_id='%s' group by  bk_id"%(bk[0])
            print(sql)
            cursor.execute(sql)
            data = cursor.fetchall()

            for dd in data:
                print(dd)
                if(dd[0] and dd[1]>2):
                    # 构建links数据
                    tmp = {}
                    name.append(bk[1])
                    name.append(dd[0])
                    tmp['source'] = bk[1]
                    tmp['target'] = dd[0]
                    tmp['value'] = dd[1]
                    re.append(tmp)

    # name去重复
    name = list(set(name))
    print(name)
    print(re)

    names = []
    for nm in name:
        nt = {}
        nt['name'] = nm
        names.append(nt)

    # print(names)
    conn.close()

    # f_data.write(json.dumps(re))
    # f_links.write(json.dumps(names))

    # 生成桑基图
    sankey = Sankey("龙虎榜营业部(玉雪锍霜)", width=1200, height=960)
    sankey.add("QQ:762285551", names, re, line_opacity=0.2,
               line_curve=0.5, line_color='source',
               is_label_show=True, label_pos='right')
    sankey.render()
